from .mae_mm_native_guide import ViTMAEConfigMM, ViTMAEForPreTrainingMM
from .vit_mm_native import ViTConfigMM, ViTModelMM


def build_model(args):
    cfg = ViTMAEConfigMM(
        mask_ratio=args.mask_ratio,
        modality_list=tuple(args.modality_list),
        num_channels=args.total_dim,
        seq_len=args.seq_len,
        embed_dim=args.embed_dim,
        depth=args.depth,
        num_heads=args.num_heads,
        decoder_embed_dim=args.decoder_embed_dim,
        decoder_depth=args.decoder_depth,
        decoder_num_heads=args.decoder_num_heads,
        norm_pix_loss=args.norm_pix_loss,
        loss_type=args.loss_type,
        guide_mode=args.guide_mode,
        guide_scale=args.guide_scale,
        num_prefix_tokens=args.num_prefix_tokens,
        use_flash_attn=args.use_flash_attn
    )
    print("build_model, cfg.modality_list: {0}, args.modality_list: {1}".format(
        cfg.modality_list, args.modality_list)
    )
    model = ViTMAEForPreTrainingMM(cfg)
    return model


def build_backbone(args):
    cfg = ViTConfigMM(
        modality_list=tuple(args.modality_list),
        num_channels=args.total_dim,
        seq_len=args.seq_len,
        embed_dim=args.embed_dim,
        depth=args.depth,
        num_heads=args.num_heads,
        drop_path_rate=args.drop_path_rate,
        num_prefix_tokens=args.num_prefix_tokens,
        use_flash_attn=args.use_flash_attn
    )
    model = ViTModelMM(cfg)
    return model
